//============================================================================
// Name        : RankBuilder.cpp
// Author      : lucmir@dcc.ufmg.br
// Description :
//============================================================================

#include "RankBuilder.h"

RankBuilder::RankBuilder(DocumentIndex *docIdx, Index *index, DocumentFeatures *docFeatures) {
	this->docIdx = docIdx;
	this->index = index;
	this->docFeatures = docFeatures;
	this->model = VECTORIAL;
	k = BM25_K;
	b = BM25_B;
}

RankBuilder::~RankBuilder() {
}

Rank *RankBuilder::generateRank(std::string *query) {
	Rank *rank = new Rank();

	std::string term;
	std::stringstream ss(*query);

	// Vectorial Model
	if(this->model == VECTORIAL) {
		while( ss >> term ) {
			calcTermRankScoreVM(rank, &term);
		}
		updateDocScoresVM(rank);
	}
	// BM25 Model
	else if(this->model == BM25) {
		while( ss >> term ) {
			calcTermRankScoreBM25(rank, &term);
		}
	}
	// Combined Model
	else if (this->model == COMBINED) {
		while (ss >> term) {
			calcTermRankScoreCM(rank, &term);
		}
	}
	// Combined with Page Rank Model
	else if (this->model == PR) {
		while (ss >> term) {
			calcTermRankScoreCM(rank, &term);
		}
		updateDocScoresPR(rank);
	}
	// Combined with Anchor Text Model
	else if (this->model == AT) {
		while (ss >> term) {
			calcTermRankScoreCMWithAT(rank, &term);
		}
	}
	// Combined with Page Rank and Anchor Text Model
	else if (this->model == PR_AT) {
		while (ss >> term) {
			calcTermRankScoreCMWithAT(rank, &term);
		}
		updateDocScoresPR(rank);
	}

	return rank;
}

/**
 * BM25 model functions
 */
void RankBuilder::calcTermRankScoreBM25(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, numDocsTfDiff, fdt, docId;
	double docLen, bdt, wdt;

	if(qr == NULL) {
		return;
	}

	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	numDocsTfDiff = numDocs - tf;

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		docLen = this->docIdx->getDoc(docId)->getLen();

		// B = [ (K + 1)*f_d,t ] / [ K*( (1-b) + b*len(d_j)/avgDocsLen ) + f_i,j ]
		bdt = ( (this->k + 1) * fdt ) / ( this->k * ( (1 - this->b) + this->b * docLen ) );

		// w_d,t = B * log( (N-n_i + 0.5) / (n_i + 0.5) )
		wdt = bdt * log( (numDocsTfDiff) / (tf + 0.5) );

		// Update accumulator
		rank->updateAccumulator(docId, wdt);
	}

	delete qr;
}

/**
 * Vectorial model functions
 */
void RankBuilder::calcTermRankScoreVM(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, fdt, docId;
	double wt, wdt;

	if(qr == NULL) {
		return;
	}

	// wt = loge(1 + N/f_t)
	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	wt = log(1 + ((double)numDocs/(double)tf) );

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		// w_d,t = loge(1 + f_d,t)*w_t
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		wdt = log(1 + (double)fdt)*wt;

		// Update accumulator
		rank->updateAccumulator(docId, wdt);
	}

	delete qr;
}

void RankBuilder::updateDocScoresVM(Rank *rank) {
	AccumulatorsItr itr;
	unsigned int docId;
	double wd, score;

	for(itr=rank->getAccumulators()->begin(); itr!=rank->getAccumulators()->end(); itr++) {
		docId = itr->first;
		wd = this->docIdx->getDoc(docId)->getWeight();
		score = (wd != 0) ? (itr->second / wd) : wd;
		rank->setAccumulator(docId, score);
	}
}

/**
 * Combined model functions
 */
void RankBuilder::calcTermRankScoreCM(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, numDocsTfDiff, fdt, docId;
	double docLen, bdt, wdt, wdt2;

	if(qr == NULL) {
		return;
	}

	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	numDocsTfDiff = numDocs - tf;
	double wt = log(1 + ((double)numDocs/(double)tf) );

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		docLen = this->docIdx->getDoc(docId)->getLen();

		// B = [ (K + 1)*f_d,t ] / [ K*( (1-b) + b*len(d_j)/avgDocsLen ) + f_i,j ]
		bdt = ( (this->k + 1) * fdt ) / ( this->k * ( (1 - this->b) + this->b * docLen ) );

		// w_d,t = B * log( (N-n_i + 0.5) / (n_i + 0.5) )
		wdt = bdt * log( (numDocsTfDiff) / (tf + 0.5) );

		wdt2 = log(1 + (double)fdt) * wt;

		// Update accumulator
		rank->updateAccumulator(docId, wdt * wdt2);
	}

	delete qr;
}

void RankBuilder::calcTermRankScoreCMWithPR(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, numDocsTfDiff, fdt, docId;
	double docLen, bdt, wdt, wdt2;

	if(qr == NULL) {
		return;
	}

	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	numDocsTfDiff = numDocs - tf;
	double wt = log(1 + ((double)numDocs/(double)tf) );

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		docLen = this->docIdx->getDoc(docId)->getLen();

		// B = [ (K + 1)*f_d,t ] / [ K*( (1-b) + b*len(d_j)/avgDocsLen ) + f_i,j ]
		bdt = ( (this->k + 1) * fdt ) / ( this->k * ( (1 - this->b) + this->b * docLen ) );

		// w_d,t = B * log( (N-n_i + 0.5) / (n_i + 0.5) )
		wdt = bdt * log( (numDocsTfDiff) / (tf + 0.5) );

		wdt2 = log(1 + (double)fdt) * wt;

		// Update accumulator
		rank->updateAccumulator(docId, wdt * wdt2);
	}

	delete qr;
}

void RankBuilder::updateDocScoresPR(Rank *rank) {
	AccumulatorsItr itr;
	unsigned int docId;
	double wd, score;

	for(itr=rank->getAccumulators()->begin(); itr!=rank->getAccumulators()->end(); itr++) {
		docId = itr->first;
		wd = this->docFeatures->getPageRank(docId);
		score = itr->second * (1.0 + wd);
		rank->setAccumulator(docId, score);
	}
}

void RankBuilder::calcTermRankScoreCMWithAT(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, numDocsTfDiff, fdt, docId, termId, anchorTextFreq;
	double docLen, bdt, wdt, wdt2, score;

	if(qr == NULL) {
		return;
	}

	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	numDocsTfDiff = numDocs - tf;
	double wt = log(1 + ((double)numDocs/(double)tf) );

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		docLen = this->docIdx->getDoc(docId)->getLen();

		// B = [ (K + 1)*f_d,t ] / [ K*( (1-b) + b*len(d_j)/avgDocsLen ) + f_i,j ]
		bdt = ( (this->k + 1) * fdt ) / ( this->k * ( (1 - this->b) + this->b * docLen ) );

		// w_d,t = B * log( (N-n_i + 0.5) / (n_i + 0.5) )
		wdt = bdt * log( (numDocsTfDiff) / (tf + 0.5) );

		wdt2 = log(1 + (double)fdt) * wt;

		// Consider anchor text
		termId = this->index->getLexon()->getTermId(*term);
		anchorTextFreq = this->docFeatures->containsAnchorText(termId, docId);
		score = (wdt + anchorTextFreq) * wdt2;

		// Update accumulator
		rank->updateAccumulator(docId, score);
	}

	// Considering Anchor Text
	/*unsigned int termId = this->index->getLexon()->getTermId();
	std::map<unsigned int, unsigned int> anchorTexts = this->docFeatures->containsAnchorText(termId);
	std::list<AnchorText::AnchorTextItem *>::iterator itr;
	for(itr=anchorTexts->begin(); itr!=anchorTexts->end(); itr++) {
		docId = (*itr)->getDocId();
		unsigned int freq = (*itr)->getFreq();
		double anchorTextScore = freq * wt;

		// Update accumulator
		rank->updateAccumulator(docId, anchorTextScore);
	}*/

	delete qr;
}

void RankBuilder::calcTermRankScoreCMWithPRAT(Rank *rank, std::string *term) {
	QueryResult *qr = this->index->queryTerm(*term);
	unsigned int numDocs, tf, numDocsTfDiff, fdt, docId;
	double docLen, bdt, wdt, wdt2;

	if(qr == NULL) {
		return;
	}

	numDocs = this->docIdx->getNumDocs();
	tf = qr->getTf();
	numDocsTfDiff = numDocs - tf;
	double wt = log(1 + ((double)numDocs/(double)tf) );

	// For each (d, f_d,t) in I_t
	unsigned int lim = tf*2;
	for(unsigned int i=0; i<lim; i+=2) {
		docId = qr->getDocFreqs()[i];
		fdt = qr->getDocFreqs()[i+1];
		docLen = this->docIdx->getDoc(docId)->getLen();

		// B = [ (K + 1)*f_d,t ] / [ K*( (1-b) + b*len(d_j)/avgDocsLen ) + f_i,j ]
		bdt = ( (this->k + 1) * fdt ) / ( this->k * ( (1 - this->b) + this->b * docLen ) );

		// w_d,t = B * log( (N-n_i + 0.5) / (n_i + 0.5) )
		wdt = bdt * log( (numDocsTfDiff) / (tf + 0.5) );

		wdt2 = log(1 + (double)fdt) * wt;

		// Update accumulator
		rank->updateAccumulator(docId, wdt * wdt2);
	}

	delete qr;
}

/**
 * Other functions
 */

void RankBuilder::printScores(Rank *rank, int rankLen) {
	AccumulatorsItr itr, maxItr;
	double maxScore;
	int count = 0;
	std::string docUrl;

	if(rank->getAccumulators()->size() == 0) {
		std::cout << "\tNenhum documento correspondente à consulta requisitada.\n";
	}

	while( (rankLen == ALL_RANK || count < rankLen) && (rank->getAccumulators()->size() > 0) ) {
		itr = rank->getAccumulators()->begin();
		maxItr = itr;
		maxScore = itr->second;
		for(;itr!=rank->getAccumulators()->end(); itr++) {
			if(itr->second > maxScore) {
				maxItr = itr;
				maxScore = itr->second;
			}
		}
		docUrl = this->docIdx->getDoc(maxItr->first)->getUrl();
		std::cout << "\t" << ++count << ") " << docUrl;
		#ifdef PRINT_SCORE
			std::cout << "\t( SCORE: " << maxItr->second << " )\n";
		#else
			std::cout << "\n";
		#endif
		rank->getAccumulators()->erase(maxItr);
	}
}


void RankBuilder::printScores(Rank *rank, int rankLen, std::string fileName) {
	AccumulatorsItr itr, maxItr;
	double maxScore;
	int count = 0;

	if(rank->getAccumulators()->size() == 0) {
		std::cout << "\tNenhum documento correspondente à consulta requisitada.\n";
	}

	std::ofstream file;
	if(this->model == VECTORIAL) {
		fileName += "_v";
	}
	else if (this->model == BM25) {
		fileName += "_b";
	}
	else if (this->model == COMBINED) {
		fileName += "_c";
	}
	else if (this->model == PR) {
		fileName += "_pr";
	}
	else if (this->model == AT) {
		fileName += "_at";
	}
	else if (this->model == PR_AT) {
		fileName += "_prat";
	}
	file.open(fileName.c_str());
	if(!file.is_open()) {
		std::cout << "\nERRO: Não foi possível abrir o arquivo: " << fileName << "\n";
	}

	while( (rankLen == ALL_RANK || count < rankLen) && (rank->getAccumulators()->size() > 0) ) {
		itr = rank->getAccumulators()->begin();
		maxItr = itr;
		maxScore = itr->second;
		for(;itr!=rank->getAccumulators()->end(); itr++) {
			if(itr->second > maxScore) {
				maxItr = itr;
				maxScore = itr->second;
			}
		}
		file << maxItr->first << "\n";
		rank->getAccumulators()->erase(maxItr);
		count++;
	}

	file.close();
}

void RankBuilder::setModel(Model model) {
	this->model = model;
}
